import subprocess
import platform
import socket
import psutil
import winreg
import os
from PyQt6.QtGui import QFont, QPainter, QColor
from PyQt6.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QGridLayout, QTableWidgetItem, QSizePolicy
from PyQt6.QtCore import Qt, QThread, pyqtSignal, QTimer
from qfluentwidgets import (CardWidget, BodyLabel, StrongBodyLabel, PushButton, FluentIcon,
                            InfoBar, InfoBarPosition, TableWidget, ScrollArea, ProgressBar,
                            PrimaryPushButton, TransparentPushButton, PillPushButton, TitleLabel)
import time
from datetime import datetime

from common.widgets.info_card import InfoCard


class ScanWorker(QThread):
    """扫描工作线程"""
    progress_updated = pyqtSignal(int, str)
    scan_completed = pyqtSignal(list, dict)  # 扫描结果, 统计数据

    def __init__(self):
        super().__init__()
        self.vulnerabilities = []
        self.stats = {}

    def run(self):
        """执行扫描任务"""
        self.vulnerabilities.clear()

        scan_steps = [
            ("检查系统信息...", self.scan_system_info),
            ("扫描开放端口...", self.scan_open_ports),
            ("检查防火墙状态...", self.scan_firewall_status),
            ("检查系统更新...", self.scan_system_updates),
            ("扫描运行服务...", self.scan_running_services),
            ("检查用户账户...", self.scan_user_accounts),
            ("检查共享资源...", self.scan_shared_resources),
            ("检查启动项...", self.scan_startup_items)
        ]

        total_steps = len(scan_steps)
        for i, (step_name, scan_func) in enumerate(scan_steps):
            self.progress_updated.emit(int((i / total_steps) * 100), step_name)
            scan_func()
            time.sleep(0.3)

        self.calculate_stats()
        self.progress_updated.emit(100, "扫描完成")
        self.scan_completed.emit(self.vulnerabilities, self.stats)

    def calculate_stats(self):
        """计算统计数据"""
        risk_counts = {"高": 0, "中": 0, "低": 0}
        type_counts = {}

        for vuln in self.vulnerabilities:
            risk = vuln.get('risk', '低')
            vuln_type = vuln.get('type', '其他')

            if risk in risk_counts:
                risk_counts[risk] += 1

            if vuln_type in type_counts:
                type_counts[vuln_type] += 1
            else:
                type_counts[vuln_type] = 1

        self.stats = {
            'total': len(self.vulnerabilities),
            'risk_counts': risk_counts,
            'type_counts': type_counts
        }

    def scan_system_info(self):
        """检查系统信息"""
        try:
            system = platform.system()
            release = platform.release()

            if system == "Windows":
                if release in ["7", "8", "8.1"]:
                    self.vulnerabilities.append({
                        'type': '系统版本',
                        'description': f'操作系统 {system} {release} 已停止支持',
                        'risk': '高',
                        'suggestion': '升级到受支持的操作系统版本',
                        'details': '过旧的操作系统不再接收安全更新'
                    })
                else:
                    self.vulnerabilities.append({
                        'type': '系统版本',
                        'description': f'操作系统 {system} {release}',
                        'risk': '低',
                        'suggestion': '保持系统更新',
                        'details': '当前系统版本正常'
                    })

            python_version = platform.python_version()
            if tuple(map(int, python_version.split('.'))) < (3, 8, 0):
                self.vulnerabilities.append({
                    'type': '软件版本',
                    'description': f'Python {python_version} 版本过旧',
                    'risk': '中',
                    'suggestion': '升级Python到3.8或更高版本',
                    'details': '旧版本可能存在安全漏洞'
                })

        except Exception as e:
            self.vulnerabilities.append({
                'type': '系统检查',
                'description': f'系统信息检查失败: {str(e)}',
                'risk': '中',
                'suggestion': '手动检查系统信息',
                'details': '无法自动获取系统信息'
            })

    def scan_open_ports(self):
        """扫描开放端口"""
        dangerous_ports = [21, 23, 135, 139, 445, 1433, 3389, 5432]
        common_ports = [22, 25, 53, 80, 110, 143, 443, 993, 995]

        for port in dangerous_ports + common_ports:
            try:
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.settimeout(0.5)
                result = sock.connect_ex(('127.0.0.1', port))

                if result == 0:
                    risk = '高' if port in dangerous_ports else '中'
                    port_name = self.get_port_name(port)

                    self.vulnerabilities.append({
                        'type': '开放端口',
                        'description': f'端口 {port} ({port_name}) 开放',
                        'risk': risk,
                        'suggestion': '关闭不必要的端口或加强访问控制',
                        'details': f'开放的{port_name}端口可能存在安全风险'
                    })
                sock.close()
            except Exception:
                pass

    def get_port_name(self, port):
        """获取端口对应的服务名"""
        port_names = {
            21: 'FTP', 22: 'SSH', 23: 'Telnet', 25: 'SMTP',
            53: 'DNS', 80: 'HTTP', 110: 'POP3', 135: 'RPC',
            139: 'NetBIOS', 143: 'IMAP', 443: 'HTTPS',
            445: 'SMB', 993: 'IMAPS', 995: 'POP3S',
            1433: 'SQL Server', 3389: 'RDP', 5432: 'PostgreSQL'
        }
        return port_names.get(port, '未知服务')

    def scan_firewall_status(self):
        """检查防火墙状态"""
        try:
            if platform.system() == 'Windows':
                result = subprocess.run(
                    ['netsh', 'advfirewall', 'show', 'allprofiles', 'state'],
                    capture_output=True, text=True, encoding='utf-8', errors='replace'
                )

                if 'OFF' in result.stdout.upper():
                    self.vulnerabilities.append({
                        'type': '防火墙',
                        'description': 'Windows防火墙已关闭',
                        'risk': '高',
                        'suggestion': '启用Windows防火墙',
                        'details': '防火墙是抵御网络攻击的第一道防线'
                    })
                else:
                    self.vulnerabilities.append({
                        'type': '防火墙',
                        'description': 'Windows防火墙已启用',
                        'risk': '低',
                        'suggestion': '保持防火墙开启状态',
                        'details': '防火墙状态正常'
                    })
        except Exception as e:
            self.vulnerabilities.append({
                'type': '防火墙',
                'description': f'防火墙状态检查失败: {str(e)}',
                'risk': '中',
                'suggestion': '手动检查防火墙状态',
                'details': '无法自动检测防火墙状态'
            })

    def scan_system_updates(self):
        """检查系统更新"""
        try:
            if platform.system() == 'Windows':
                result = subprocess.run(
                    ['sc', 'query', 'wuauserv'],
                    capture_output=True, text=True, encoding='utf-8', errors='replace'
                )

                if 'STOPPED' in result.stdout.upper():
                    self.vulnerabilities.append({
                        'type': '系统更新',
                        'description': 'Windows更新服务已停止',
                        'risk': '中',
                        'suggestion': '启用Windows更新服务',
                        'details': '自动更新有助于及时获取安全补丁'
                    })
                else:
                    self.vulnerabilities.append({
                        'type': '系统更新',
                        'description': 'Windows更新服务正在运行',
                        'risk': '低',
                        'suggestion': '定期检查并安装更新',
                        'details': '更新服务状态正常'
                    })
        except Exception as e:
            self.vulnerabilities.append({
                'type': '系统更新',
                'description': f'更新状态检查失败: {str(e)}',
                'risk': '中',
                'suggestion': '手动检查系统更新设置',
                'details': '无法自动检测更新状态'
            })

    def scan_running_services(self):
        """扫描运行中的服务"""
        try:
            dangerous_services = ['telnet', 'ftp', 'rsh', 'rlogin']

            for proc in psutil.process_iter(['pid', 'name']):
                try:
                    if proc.info['name'] and any(svc in proc.info['name'].lower() for svc in dangerous_services):
                        self.vulnerabilities.append({
                            'type': '危险服务',
                            'description': f'发现运行中的服务: {proc.info["name"]}',
                            'risk': '高',
                            'suggestion': '停止不安全的服务',
                            'details': '此服务可能存在安全风险'
                        })
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue

        except Exception as e:
            self.vulnerabilities.append({
                'type': '服务扫描',
                'description': f'服务扫描失败: {str(e)}',
                'risk': '低',
                'suggestion': '手动检查运行中的服务',
                'details': '建议定期审查系统服务'
            })

    def scan_user_accounts(self):
        """检查用户账户"""
        try:
            if platform.system() == 'Windows':
                result = subprocess.run(
                    ['net', 'user', 'guest'],
                    capture_output=True, text=True, encoding='utf-8', errors='replace'
                )

                if 'Account active' in result.stdout and 'Yes' in result.stdout:
                    self.vulnerabilities.append({
                        'type': '用户账户',
                        'description': 'Guest账户已启用',
                        'risk': '高',
                        'suggestion': '禁用Guest账户',
                        'details': 'Guest账户可能被恶意利用'
                    })

                result = subprocess.run(
                    ['net', 'localgroup', 'administrators'],
                    capture_output=True, text=True, encoding='utf-8', errors='replace'
                )

                admin_count = len(
                    [line for line in result.stdout.split('\n') if line.strip() and not line.startswith('-')])
                if admin_count > 5:
                    self.vulnerabilities.append({
                        'type': '用户账户',
                        'description': '管理员账户过多',
                        'risk': '中',
                        'suggestion': '减少不必要的管理员权限',
                        'details': '过多的管理员账户增加安全风险'
                    })

        except Exception as e:
            self.vulnerabilities.append({
                'type': '用户账户',
                'description': f'用户账户检查失败: {str(e)}',
                'risk': '低',
                'suggestion': '手动检查用户账户设置',
                'details': '建议定期审查用户权限'
            })

    def scan_shared_resources(self):
        """检查共享资源"""
        try:
            if platform.system() == 'Windows':
                result = subprocess.run(
                    ['net', 'share'],
                    capture_output=True, text=True, encoding='utf-8', errors='replace'
                )

                shares = [line for line in result.stdout.split('\n')
                          if line.strip() and not line.startswith('-') and '$' not in line]

                if len(shares) > 3:
                    self.vulnerabilities.append({
                        'type': '共享资源',
                        'description': f'发现 {len(shares)} 个共享文件夹',
                        'risk': '中',
                        'suggestion': '检查共享权限，关闭不必要的共享',
                        'details': '不当的文件共享可能导致数据泄露'
                    })

        except Exception as e:
            self.vulnerabilities.append({
                'type': '共享资源',
                'description': f'共享资源检查失败: {str(e)}',
                'risk': '低',
                'suggestion': '手动检查网络共享设置',
                'details': '建议定期审查共享权限'
            })

    def scan_startup_items(self):
        """检查启动项"""
        try:
            if platform.system() == 'Windows':
                startup_locations = [
                    r"SOFTWARE\Microsoft\Windows\CurrentVersion\Run",
                    r"SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce"
                ]

                startup_count = 0
                for location in startup_locations:
                    try:
                        key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, location)
                        i = 0
                        while True:
                            try:
                                winreg.EnumValue(key, i)
                                startup_count += 1
                                i += 1
                            except WindowsError:
                                break
                        winreg.CloseKey(key)
                    except FileNotFoundError:
                        continue

                if startup_count > 10:
                    self.vulnerabilities.append({
                        'type': '启动项',
                        'description': f'发现 {startup_count} 个启动项',
                        'risk': '中',
                        'suggestion': '清理不必要的启动项',
                        'details': '过多启动项影响系统性能和安全'
                    })

        except Exception as e:
            self.vulnerabilities.append({
                'type': '启动项',
                'description': f'启动项检查失败: {str(e)}',
                'risk': '低',
                'suggestion': '手动检查系统启动项',
                'details': '建议定期清理启动项'
            })


class Scanner(QWidget):
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.vulnerabilities = []
        self.stats = {}
        self.scan_worker = None

    def init_ui(self):
        # 主布局
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(30, 30, 30, 30)
        main_layout.setSpacing(20)
        self.setFont(QFont("Microsoft YaHei", 18))

        # 标题区域
        title_layout = QHBoxLayout()
        title = TitleLabel("安全漏洞扫描")
        title_layout.addWidget(title)
        title_layout.addStretch()

        self.scan_button = PushButton("开始全面扫描", icon=FluentIcon.SEARCH)
        self.scan_button.clicked.connect(self.start_scan)
        title_layout.addWidget(self.scan_button)

        main_layout.addLayout(title_layout)

        # 统计卡片区域 - 修改布局使其居中
        stats_container = QWidget()
        stats_container_layout = QHBoxLayout(stats_container)
        stats_container_layout.setContentsMargins(0, 0, 0, 0)

        # 添加左侧弹性空间
        stats_container_layout.addStretch()

        self.total_card = InfoCard("总问题", "扫描中", "个", "📊", "#1890ff")
        self.high_card = InfoCard("高风险", "扫描中", "个", "🔴", "#ff4d4f")
        self.medium_card = InfoCard("中风险", "扫描中", "个", "🟡", "#faad14")
        self.low_card = InfoCard("低风险", "扫描中", "个", "🟢", "#52c41a")

        # 为初始状态设置灰色样式
        self.set_card_inactive_style()

        stats_container_layout.addWidget(self.total_card)
        stats_container_layout.addSpacing(15)
        stats_container_layout.addWidget(self.high_card)
        stats_container_layout.addSpacing(15)
        stats_container_layout.addWidget(self.medium_card)
        stats_container_layout.addSpacing(15)
        stats_container_layout.addWidget(self.low_card)

        # 添加右侧弹性空间
        stats_container_layout.addStretch()

        main_layout.addWidget(stats_container)

        # 进度区域
        progress_card = CardWidget()
        progress_layout = QVBoxLayout(progress_card)
        progress_layout.setContentsMargins(20, 15, 20, 15)
        progress_layout.setSpacing(10)

        self.progress_bar = ProgressBar()
        self.progress_bar.setValue(0)
        self.progress_bar.setEnabled(False)  # 非扫描状态下禁用
        self.progress_label = BodyLabel("准备就绪")
        self.progress_label.setStyleSheet("color: #666;")

        progress_layout.addWidget(self.progress_bar)
        progress_layout.addWidget(self.progress_label)

        main_layout.addWidget(progress_card)

        # 结果表格
        result_card = CardWidget()
        result_layout = QVBoxLayout(result_card)
        result_layout.setContentsMargins(20, 20, 20, 20)
        result_layout.setSpacing(15)

        table_header = QHBoxLayout()
        table_title = StrongBodyLabel("扫描结果详情")
        table_title.setFont(QFont("Microsoft YaHei", 14))
        table_header.addWidget(table_title)
        table_header.addStretch()

        self.export_button = TransparentPushButton("导出报告", icon=FluentIcon.SAVE)
        self.export_button.clicked.connect(self.export_report)
        self.export_button.setVisible(False)
        table_header.addWidget(self.export_button)

        result_layout.addLayout(table_header)

        self.result_table = TableWidget()
        self.result_table.setColumnCount(5)
        self.result_table.setHorizontalHeaderLabels(["类型", "问题描述", "风险等级", "建议措施", "详细信息"])

        header = self.result_table.horizontalHeader()
        header.setSectionResizeMode(0, header.ResizeMode.Fixed)
        header.setSectionResizeMode(1, header.ResizeMode.Stretch)
        header.setSectionResizeMode(2, header.ResizeMode.Fixed)
        header.setSectionResizeMode(3, header.ResizeMode.Stretch)
        header.setSectionResizeMode(4, header.ResizeMode.Stretch)

        self.result_table.setColumnWidth(0, 100)
        self.result_table.setColumnWidth(2, 80)

        result_layout.addWidget(self.result_table)
        main_layout.addWidget(result_card)

    def set_card_inactive_style(self):
        """为未扫描状态的卡片设置灰色样式"""
        inactive_color = "#8c8c8c"
        self.total_card.update_value("未扫描", "")
        self.high_card.update_value("未扫描", "")
        self.medium_card.update_value("未扫描", "")
        self.low_card.update_value("未扫描", "")

    def start_scan(self):
        """开始扫描"""
        self.vulnerabilities.clear()
        self.result_table.setRowCount(0)
        self.scan_button.setEnabled(False)
        self.export_button.setVisible(False)

        # 重置统计卡片为扫描中状态
        self.set_card_scanning_style()

        # 更新进度显示
        self.progress_bar.setEnabled(True)
        self.progress_bar.setValue(0)
        self.progress_label.setText("正在初始化扫描...")

        # 启动扫描线程
        self.scan_worker = ScanWorker()
        self.scan_worker.progress_updated.connect(self.update_progress)
        self.scan_worker.scan_completed.connect(self.on_scan_completed)
        self.scan_worker.start()

    def set_card_scanning_style(self):
        """为扫描中状态的卡片设置样式"""
        scanning_color = "#1890ff"
        self.total_card.update_value("扫描中...", "")
        self.high_card.update_value("扫描中...", "")
        self.medium_card.update_value("扫描中...", "")
        self.low_card.update_value("扫描中...", "")

    def update_progress(self, value, current_task):
        """更新进度"""
        self.progress_bar.setValue(value)
        self.progress_label.setText(current_task)

    def update_stats_display(self, stats):
        """更新统计显示"""
        self.total_card.update_value(str(stats.get('total', 0)), "个")
        risk_counts = stats.get('risk_counts', {})
        self.high_card.update_value(str(risk_counts.get('高', 0)), "个",)
        self.medium_card.update_value(str(risk_counts.get('中', 0)), "个")
        self.low_card.update_value(str(risk_counts.get('低', 0)), "个")

    def on_scan_completed(self, vulnerabilities, stats):
        """扫描完成"""
        self.vulnerabilities = vulnerabilities
        self.stats = stats

        # 更新统计显示
        self.update_stats_display(stats)

        # 更新表格
        self.result_table.setRowCount(len(vulnerabilities))
        for i, vuln in enumerate(vulnerabilities):
            self.result_table.setItem(i, 0, QTableWidgetItem(vuln.get('type', '')))
            self.result_table.setItem(i, 1, QTableWidgetItem(vuln.get('description', '')))

            risk_item = QTableWidgetItem(vuln.get('risk', ''))
            risk_color = {"高": "#ff4d4f", "中": "#faad14", "低": "#52c41a"}.get(vuln.get('risk', ''), "#666")
            risk_item.setForeground(QColor(risk_color))
            self.result_table.setItem(i, 2, risk_item)

            self.result_table.setItem(i, 3, QTableWidgetItem(vuln.get('suggestion', '')))
            self.result_table.setItem(i, 4, QTableWidgetItem(vuln.get('details', '')))

        # 恢复控件状态 - 修改进度条显示
        self.scan_button.setEnabled(True)
        self.progress_bar.setEnabled(True)  # 保持启用状态以显示完成状态
        self.progress_bar.setValue(100)  # 保持100%进度
        self.progress_label.setText("扫描完成")  # 显示完成状态
        self.export_button.setVisible(True)

        # 显示完成提示
        InfoBar.success(
            title="扫描完成",
            content=f"共发现 {len(vulnerabilities)} 个安全问题，其中高风险 {stats['risk_counts']['高']} 个",
            orient=Qt.Orientation.Horizontal,
            isClosable=True,
            position=InfoBarPosition.TOP_RIGHT,
            duration=3000,
            parent=self
        )

    def export_report(self):
        """导出报告"""
        try:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"security_scan_report_{timestamp}.txt"

            with open(filename, 'w', encoding='utf-8') as f:
                f.write("安全漏洞扫描报告\n")
                f.write("=" * 50 + "\n")
                f.write(f"扫描时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"扫描结果统计:\n")
                f.write(f"  总问题数: {self.stats['total']}\n")
                f.write(f"  高风险: {self.stats['risk_counts']['高']}\n")
                f.write(f"  中风险: {self.stats['risk_counts']['中']}\n")
                f.write(f"  低风险: {self.stats['risk_counts']['低']}\n\n")

                f.write("详细问题列表:\n")
                f.write("-" * 50 + "\n")
                for i, vuln in enumerate(self.vulnerabilities, 1):
                    f.write(f"{i}. [{vuln['type']}] {vuln['description']}\n")
                    f.write(f"   风险等级: {vuln['risk']}\n")
                    f.write(f"   建议措施: {vuln['suggestion']}\n")
                    f.write(f"   详细信息: {vuln['details']}\n\n")

            InfoBar.success(
                title="导出成功",
                content=f"报告已保存为 {filename}",
                orient=Qt.Orientation.Horizontal,
                isClosable=True,
                position=InfoBarPosition.TOP_RIGHT,
                duration=2000,
                parent=self
            )

        except Exception as e:
            InfoBar.error(
                title="导出失败",
                content=f"导出报告时发生错误: {str(e)}",
                orient=Qt.Orientation.Horizontal,
                isClosable=True,
                position=InfoBarPosition.TOP_RIGHT,
                duration=3000,
                parent=self
            )